Implement ONNX Pad Operator #2007

johnhuichen · 2024-07-11T17:07:50Z

Pull Request Template

Checklist

Confirmed that run-checks all script has been executed.
Made sure the book is up to date with changes in this PR.

Related Issues/PRs

#1714

Changes

Implement ONNX pad by calling existing tensor.pad

Input tensor: currently supporting only Float32 but will extend the support to other floats, int and boolean values

Pads: currently supporting only positive pads

Constant value: type restricted to the input tensor element type, currently supporting only Float32

mode: currently supporting only "constant"

TODO:

Support int, boolean tensors

Currently passing a negative constant value may cause an issue. Need to investigate further

Testing

Test ONNX model produces expected outputs for f32 input tensor and positive pads before building ONNX model
Test pad codegen creates expected rust code
Test generated pad rust code produces expected outputs for f32 input tensor and positive pads

codecov · 2024-07-11T18:19:20Z

Codecov Report

Attention: Patch coverage is 99.37500% with 1 line in your changes missing coverage. Please review.

Project coverage is 86.45%. Comparing base (35345de) to head (706b48b).
Report is 18 commits behind head on main.

Files	Patch %	Lines
crates/burn-import/src/burn/node/pad.rs	98.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2007      +/-   ##
==========================================
+ Coverage   84.40%   86.45%   +2.04%     
==========================================
  Files         842      684     -158     
  Lines      105132    87413   -17719     
==========================================
- Hits        88736    75572   -13164     
+ Misses      16396    11841    -4555

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

antimora

Thank you for your addition!

Overall looks good. However, we need to tighten the pad configuration. We need to panic aggressively if the user is passing more configuration that our current APIs support. It will be an opportunity to enhance Burn's pad function.

crates/burn-import/src/onnx/op_configuration.rs

antimora · 2024-07-15T00:06:01Z

crates/burn-import/src/onnx/op_configuration.rs

+            let left = pads[input_dim - 1];
+            let top = pads[input_dim - 2];
+            let right = pads[pads.len() - 1];
+            let bottom = pads[pads.len() - 2];
+            vec![left, right, top, bottom]


This assumes the user wants to pad last 2 dimensions which works for the Burn's pad but it's silent. We should break if ONNX configuration does not fit Burn's API. Obviously ONNX version is more generic, which requires Burn's code be updated as well (not required for this PR).

antimora · 2024-07-15T00:07:14Z

crates/burn-import/src/onnx/op_configuration.rs

+    // TODO: support int, boolean
+    let constant_value: f32 = match &node.inputs[2].value {
+        Some(Data::Float32s(shape)) => shape.first().unwrap().to_owned(),
+        _ => 0.,
+    };


We need to panic and not be silent. Also, there could be only one input (data) and attributes for others see OpSet 1

laggui

Thanks for adding this implementation! I think overall it looks good.

I agree with the comments made by @antimora, we need to be more strict/clear about the unsupported configurations.

crates/burn-import/src/onnx/op_configuration.rs

johnhuichen · 2024-07-15T14:35:53Z

@antimora @laggui

I have made onnx conversion stricter by forcing three inputs. Also added a warning about pad applying to only the last two dimensions.

can you guys take another look?

laggui

That's better! I think there's still just one issue with the current way the padding is parsed.. we should not simply warn when the operation might be incorrect.

We should make sure that the provided values will work for the Burn conversion, otherwise panic.

laggui · 2024-07-15T14:49:36Z

crates/burn-import/src/onnx/op_configuration.rs

+        if input_dim > 2 {
+            log::warn!("Pad: padding will only be applied to the last two dimensions");
+        }
+
+        let left = pads[input_dim - 1];
+        let top = pads[input_dim - 2];
+        let right = pads[pads.len() - 1];
+        let bottom = pads[pads.len() - 2];
+        vec![left, right, top, bottom]


I think the correct conversion when parsing pads should check the pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...], and we should make sure that all dimensions outside of the last two are not padded otherwise the conversion is not valid (and thus, panic instead of warn).

Following the given spec example, a 4d input with padded on the last two dimensions should have pads = [0, 0, x2_begin, x3_begin, 0, 0, x2_end, x3_end] (i.e., x1_begin, x1_end, x2_begin, and x2_end are all zero)

ok that's a good idea

I have pushed the update

antimora · 2024-07-16T02:17:52Z

crates/burn-import/src/onnx/op_configuration.rs

+
+        let left_index = input_dim - 1;
+        let top_index = input_dim - 2;
+        let right_index = pads.len() - 1;
+        let bottom_index = pads.len() - 2;
+        let index_list = [left_index, top_index, right_index, bottom_index];



This seems to accept pads with pads.len() > 4 but Burn API only accepts 4. Right?

If pads.len() > 4, this code checks if the user wants to pad any dimension than the last two:

yes -> panic
no -> the pad operation is doable, and it calculates l, t, r, b

The way the padding is specified for the ONNX Pad op is a bit different since it supports padding more than just the last two dimensions. So we have to parse the provided pads to make sure it is valid, and then return the Burn valid padding with only 4 values for (left, right, top, bottom).

I think the check here is valid.

For a 4d input where we should only accept padding the last two dimensions, it means that pads should be [0, 0, x2_begin, x3_begin, 0, 0, x2_end, x3_end]. For the last two dimensions expected to be [H, W], this means that x2_begin = top, x3_begin = left, x2_end = bottom and x3_end = right.

And the op config returns only the (left, right, top, bottom) to generate the burn code.

crates/burn-import/src/onnx/op_configuration.rs

antimora

One minor fix is required. Please see my inlined comment. We should be able to merge afterwards.

antimora · 2024-07-22T18:21:59Z

crates/burn-import/src/onnx/op_configuration.rs

+        if node.inputs.len() != 3 {
+            panic!("Pad: must provide three inputs")
+        }


The says at least 2 required (so check should be node.inputs.len() > 1)

3rd and 4th inputs are optional.

antimora · 2024-07-22T18:24:11Z

crates/burn-import/src/onnx/op_configuration.rs

+        match &node.inputs[2].value {
+            Some(Data::Float32s(shape)) => shape.first().unwrap().to_owned(),
+            _ => panic!("Pad: should provide a constant value input to pad with, for example 0.0"),
+        }


3rd input is optional and should be default to 0 if absent.

pad now requires 2 or more arguments if the third argument is not given, it will default to 0

antimora

Looks good. Thanks for fixing. Minor wording update needed but I am approving in advance.

antimora · 2024-07-23T17:45:09Z

crates/burn-import/src/onnx/op_configuration.rs

-        if node.inputs.len() != 3 {
-            panic!("Pad: must provide three inputs")
+        if node.inputs.len() < 2 {
+            panic!("Pad: must provide two inputs")


I would add "at least".

"must provide two inputs" => "must provide at least two inputs"

Change panic comment from needing two inputs. This comes from the fact that the ONNX spec requires two necessary inputs but could have more two more optional argument.

johnhuichen force-pushed the implement-onnx-pad branch from cb4cb1d to a7dc181 Compare July 11, 2024 17:08

johnhuichen changed the title ~~Create ONNX test model builder~~ Implement ONNX Pad Operator Jul 11, 2024

johnhuichen marked this pull request as ready for review July 11, 2024 17:13

johnhuichen force-pushed the implement-onnx-pad branch from 3ce48de to 02b5260 Compare July 11, 2024 17:54

antimora requested review from laggui and antimora July 12, 2024 20:03

antimora requested changes Jul 15, 2024

View reviewed changes

laggui reviewed Jul 15, 2024

View reviewed changes

crates/burn-import/src/onnx/op_configuration.rs Outdated Show resolved Hide resolved

johnhuichen force-pushed the implement-onnx-pad branch 2 times, most recently from df2fee2 to 40f2a78 Compare July 15, 2024 14:29

laggui reviewed Jul 15, 2024

View reviewed changes

johnhuichen force-pushed the implement-onnx-pad branch from d28ac8e to 9391a5d Compare July 15, 2024 15:32

antimora reviewed Jul 16, 2024

View reviewed changes

laggui reviewed Jul 17, 2024

View reviewed changes

crates/burn-import/src/onnx/op_configuration.rs Outdated Show resolved Hide resolved

johnhuichen force-pushed the implement-onnx-pad branch from 5e26597 to af36c6d Compare July 19, 2024 12:39

Implement ONNX pad

2859976

johnhuichen force-pushed the implement-onnx-pad branch from af36c6d to 2859976 Compare July 19, 2024 13:02

laggui approved these changes Jul 19, 2024

View reviewed changes

laggui requested a review from antimora July 22, 2024 12:39

antimora requested changes Jul 22, 2024

View reviewed changes

mepatrick73 added 2 commits July 23, 2024 11:58

ONNX pad arguments fix

c7b0029

pad now requires 2 or more arguments if the third argument is not given, it will default to 0

fixing bug in input len fix

706b48b

mepatrick73 requested a review from antimora July 23, 2024 17:24

antimora approved these changes Jul 23, 2024

View reviewed changes

change panic comment

f230da1

Change panic comment from needing two inputs. This comes from the fact that the ONNX spec requires two necessary inputs but could have more two more optional argument.

mepatrick73 merged commit 4a3fc9d into tracel-ai:main Jul 23, 2024
6 of 12 checks passed

johnhuichen deleted the implement-onnx-pad branch July 24, 2024 16:08

antimora mentioned this pull request Aug 29, 2024

Update SUPPORTED-ONNX-OPS.md #2217

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement ONNX Pad Operator #2007

Implement ONNX Pad Operator #2007

johnhuichen commented Jul 11, 2024 •

edited

Loading

codecov bot commented Jul 11, 2024 •

edited

Loading

antimora left a comment

antimora Jul 15, 2024

antimora Jul 15, 2024

laggui left a comment

johnhuichen commented Jul 15, 2024

laggui left a comment

laggui Jul 15, 2024

johnhuichen Jul 15, 2024

johnhuichen Jul 15, 2024

antimora Jul 16, 2024

johnhuichen Jul 16, 2024

laggui Jul 17, 2024 •

edited

Loading

antimora left a comment

antimora Jul 22, 2024

antimora Jul 22, 2024

antimora left a comment

antimora Jul 23, 2024

Implement ONNX Pad Operator #2007

Implement ONNX Pad Operator #2007

Conversation

johnhuichen commented Jul 11, 2024 • edited Loading

Pull Request Template

Checklist

Related Issues/PRs

Changes

Testing

codecov bot commented Jul 11, 2024 • edited Loading

Codecov Report

antimora left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

laggui left a comment

Choose a reason for hiding this comment

johnhuichen commented Jul 15, 2024

laggui left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

laggui Jul 17, 2024 • edited Loading

Choose a reason for hiding this comment

antimora left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

antimora left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

johnhuichen commented Jul 11, 2024 •

edited

Loading

codecov bot commented Jul 11, 2024 •

edited

Loading

laggui Jul 17, 2024 •

edited

Loading